﻿using System.Runtime.CompilerServices;

namespace LiteVectorIndex.Engine
{
    /// <summary>
    /// 轻量级向量集合
    /// </summary>
    public class LiteVectorCollection
    {
        public LiteVectorCollection() : this(MetricType.Cosine)
        {
        }

        public LiteVectorCollection(MetricType metricType)
        {
            this._metricType = metricType;
        }

        public const int DEFAULT_DIMENSION = 1024; // 默认维度1024
        public const int STRUCTRE_VERSION = 3; // 当前数据结构的版本号

        private readonly List<float[]> _vectors = new(); // 向量数组
        private readonly List<int> _ids = new(); // 向量对应的ID数组，顺序要严格匹配

        private MetricType _metricType;

        /// <summary>
        /// 维度数量
        /// </summary>
        public int Dimension
        {
            get
            {
                if (_vectors.Count > 0) return _vectors[0].Length;
                else return DEFAULT_DIMENSION;
            }
        }

        /// <summary>
        /// 集合采用的距离计算方式
        /// </summary>
        public MetricType MetricTypeInUse
        {
            get
            {
                return _metricType;
            }
        }

        /// <summary>
        /// 是否做归一化处理
        /// <summary>
        public bool NormalizeOnAddAndQuery
        {
            get { return _metricType == MetricType.Cosine; }
        }

        /// <summary>
        /// 添加向量
        /// </summary>
        /// <param name="id">向量的ID</param>
        /// <param name="vector">向量本身</param>
        /// <exception cref="ArgumentException"></exception>
        public void AddVector(int id, float[] vector)
        {
            if (_vectors.Count > 0 && vector.Length != Dimension)
            {
                throw new ArgumentException($"当前向量的维度 {vector.Length} 与集合中已存在向量的维度 {Dimension} 不匹配。");
            }

            RemoveVector(id); // 确保ID唯一性，先删除旧的向量（如果存在）

            _ids.Add(id);
            _vectors.Add(vector);
        }

        /// <summary>
        /// 批量添加向量
        /// </summary>
        /// <param name="items">向量数组</param>
        public void AddVectors(IEnumerable<(int id, float[] Vector)> items)
        {
            foreach (var (id, vector) in items)
            {
                AddVector(id, vector);
            }
        }

        /// <summary>
        /// 删除向量
        /// </summary>
        /// <param name="id">被删除向量的ID，不存在的话，就跳过</param>
        /// <returns>是否真的完成删除，false意味着集合中本来就没有这个ID</returns>
        public bool RemoveVector(int id)
        {
            int index = _ids.IndexOf(id);

            if (index >= 0 && index < _vectors.Count)
            {
                _ids.RemoveAt(index);
                _vectors.RemoveAt(index);

                return true;
            }
            else
            {
                return false;
            }
        }

        /// <summary>
        /// 批量删除向量
        /// </summary>
        /// <param name="ids"></param>
        /// <returns>实际被删除的ID列表</returns>
        public IEnumerable<int> RemoveVectors(IEnumerable<int> ids)
        {
            List<int> result = [];

            foreach (var id in ids)
            {
                if (RemoveVector(id))
                {
                    result.Add(id);
                }
            }

            return result;
        }

        /// <summary>
        /// 按照近似度检索
        /// </summary>
        /// <param name="query">查询的向量，维度不匹配时，直接返回空</param>
        /// <param name="topK">返回不超过k个元素</param>
        /// <returns></returns>
        public IEnumerable<SearchResult> Search(float[] query, int topK = 3)
        {
            // 维度不匹配，不会返回任何数据。
            if (query.Length != Dimension)
            {
                return [];
            }


            var results = new List<SearchResult>(_vectors.Count);

            for (int i = 0; i < _vectors.Count; i++)
            {
                float d = DistanceProvider.Similarity(query, _vectors[i], MetricTypeInUse);
                results.Add(new SearchResult(_ids[i], d));
            }

            return [.. results.OrderByDescending(r => r.score).Take(topK)];


        }

        /// <summary>
        /// 清空ID和向量两个列表
        /// </summary>
        public void Clear()
        {
            this._ids.Clear();
            this._vectors.Clear();
        }

        /// <summary>
        /// 提取集合的配置和数据
        /// </summary>
        /// <returns>数据包</returns>
        public PlainCollectionData GetData()
        {
            var data = new PlainCollectionData
            {
                Version = STRUCTRE_VERSION,
                MetricTypeInUse = MetricTypeInUse,
                Vectors = _vectors,
                Ids = _ids
            };
            return data;
        }

        /// <summary>
        /// 使用数据包初始化当前的集合
        /// </summary>
        /// <param name="data">数据包</param>
        /// <exception cref="NotSupportedException"></exception>
        public void LoadData(PlainCollectionData data)
        {
            if (data.Version != STRUCTRE_VERSION)
                throw new NotSupportedException("不支持加载旧版本的数据结构，请重新创建向量集合并导入数据。");

            this.Clear();

            _metricType = data.MetricTypeInUse;

            this._ids.AddRange(data.Ids);
            this._vectors.AddRange(data.Vectors);
        }

        /// <summary>
        /// 集合中的向量数量
        /// </summary>
        public long Count => _vectors.Count;

    }


}
